package main

import (
	"bytes"
	"crypto/cipher"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"github.com/tjfoc/gmsm/sm4"
	"io"
	"log"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

// RateLimiter 滑动窗口限流器
type RateLimiter struct {
	windowSize  time.Duration // 时间窗口大小
	maxRequests int           // 最大请求数
	requests    []time.Time   // 请求时间记录
	mu          sync.RWMutex
}

// NewRateLimiter 创建新的限流器
func NewRateLimiter(windowSize time.Duration, maxRequests int) *RateLimiter {
	return &RateLimiter{
		windowSize:  windowSize,
		maxRequests: maxRequests,
		requests:    make([]time.Time, 0),
	}
}

// Allow 检查是否允许请求
func (rl *RateLimiter) Allow() bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	windowStart := now.Add(-rl.windowSize)

	// 清理过期的请求记录
	validRequests := make([]time.Time, 0)
	for _, reqTime := range rl.requests {
		if reqTime.After(windowStart) {
			validRequests = append(validRequests, reqTime)
		}
	}
	rl.requests = validRequests

	// 检查是否超过限制
	if len(rl.requests) >= rl.maxRequests {
		return false
	}

	// 记录当前请求
	rl.requests = append(rl.requests, now)
	return true
}

// GetStats 获取限流统计信息
func (rl *RateLimiter) GetStats() (int, int) {
	rl.mu.RLock()
	defer rl.mu.RUnlock()

	now := time.Now()
	windowStart := now.Add(-rl.windowSize)

	currentRequests := 0
	for _, reqTime := range rl.requests {
		if reqTime.After(windowStart) {
			currentRequests++
		}
	}

	return currentRequests, rl.maxRequests
}

// URLRateLimiter URL路径限流管理器
type URLRateLimiter struct {
	limiters    map[string]*RateLimiter // URL路径 -> 限流器映射
	windowSize  time.Duration           // 时间窗口大小
	maxRequests int                     // 最大请求数
	mu          sync.RWMutex            // 读写锁
}

// NewURLRateLimiter 创建新的URL限流管理器
func NewURLRateLimiter(windowSize time.Duration, maxRequests int) *URLRateLimiter {
	return &URLRateLimiter{
		limiters:    make(map[string]*RateLimiter),
		windowSize:  windowSize,
		maxRequests: maxRequests,
	}
}

// Allow 检查指定URL路径是否允许请求
func (urlrl *URLRateLimiter) Allow(urlPath string) bool {
	// 获取或创建URL路径对应的限流器
	limiter := urlrl.getOrCreateLimiter(urlPath)
	return limiter.Allow()
}

// GetStats 获取指定URL路径的限流统计信息
func (urlrl *URLRateLimiter) GetStats(urlPath string) (int, int) {
	limiter := urlrl.getOrCreateLimiter(urlPath)
	return limiter.GetStats()
}

// GetAllStats 获取所有URL路径的限流统计信息
func (urlrl *URLRateLimiter) GetAllStats() map[string]map[string]interface{} {
	urlrl.mu.RLock()
	defer urlrl.mu.RUnlock()

	stats := make(map[string]map[string]interface{})
	for urlPath, limiter := range urlrl.limiters {
		current, max := limiter.GetStats()
		stats[urlPath] = map[string]interface{}{
			"current_requests": current,
			"max_requests":     max,
			"remaining":        max - current,
			"window_size":      urlrl.windowSize.Seconds(),
		}
	}
	return stats
}

// getOrCreateLimiter 获取或创建URL路径对应的限流器
func (urlrl *URLRateLimiter) getOrCreateLimiter(urlPath string) *RateLimiter {
	urlrl.mu.Lock()
	defer urlrl.mu.Unlock()

	// 如果限流器不存在，则创建新的
	if limiter, exists := urlrl.limiters[urlPath]; exists {
		return limiter
	}

	// 创建新的限流器
	limiter := NewRateLimiter(urlrl.windowSize, urlrl.maxRequests)
	urlrl.limiters[urlPath] = limiter
	return limiter
}

// cleanup 清理长时间未使用的限流器（可选的内存优化）
func (urlrl *URLRateLimiter) cleanup() {
	urlrl.mu.Lock()
	defer urlrl.mu.Unlock()

	// 这里可以添加清理逻辑，比如删除长时间未使用的限流器
	// 为了简化，暂时不实现清理逻辑
}

type Proxy struct {
	requestBlock   cipher.Block // 用于解密请求的加密块
	responseBlock  cipher.Block // 用于加密响应的加密块
	requestKey     []byte       // 请求解密密钥
	responseKey    []byte       // 响应加密密钥
	address        string
	dataKey        string
	client         *http.Client
	mu             sync.RWMutex
	urlRateLimiter *URLRateLimiter // URL路径限流管理器
}

func main() {
	var (
		port           string
		requestKeyStr  string
		responseKeyStr string
		targetAddr     string
		dataKey        string
		timeout        int
		rateLimit      string
	)

	// 定义标志和默认值
	flag.StringVar(&port, "port", "18090", "代理服务器监听的端口")
	flag.StringVar(&requestKeyStr, "request-key", "1809012345678910", "请求解密SM4密钥")
	flag.StringVar(&responseKeyStr, "response-key", "1809012345678910", "响应加密SM4密钥")
	flag.StringVar(&targetAddr, "target", "127.0.0.1:81", "目标服务器地址")
	flag.StringVar(&dataKey, "data-key", "data", "数据键名")
	flag.IntVar(&timeout, "timeout", 30, "HTTP请求超时时间(秒)")
	flag.StringVar(&rateLimit, "rate-limit", "60,5", "限流配置，格式：时间窗口秒数,最大请求数，如：60,10 表示60秒内最多10次请求")

	// 解析命令行参数
	flag.Parse()

	// 验证SM4密钥长度
	if len(requestKeyStr) != 16 {
		log.Fatal("请求解密密钥长度必须为16字节")
	}
	if len(responseKeyStr) != 16 {
		log.Fatal("响应加密密钥长度必须为16字节")
	}

	// 校验和修正目标地址格式
	targetAddr = normalizeTargetAddress(targetAddr)

	// 校验URL合法性
	u, err := url.Parse(targetAddr)
	if err != nil {
		log.Fatal("无效的目标地址:", err)
	}

	// 创建请求解密SM4加密块
	requestBlock, err := sm4.NewCipher([]byte(requestKeyStr))
	if err != nil {
		log.Fatal("无效的请求解密密钥:", err)
	}

	// 创建响应加密SM4加密块
	responseBlock, err := sm4.NewCipher([]byte(responseKeyStr))
	if err != nil {
		log.Fatal("无效的响应加密密钥:", err)
	}

	// 创建自定义HTTP客户端
	client := &http.Client{
		Timeout: time.Duration(timeout) * time.Second,
		Transport: &http.Transport{
			MaxIdleConns:        100,
			MaxIdleConnsPerHost: 10,
			IdleConnTimeout:     90 * time.Second,
		},
	}

	// 解析限流配置
	var urlRateLimiter *URLRateLimiter
	if rateLimit != "" {
		parts := strings.Split(rateLimit, ",")
		if len(parts) != 2 {
			log.Fatal("限流配置格式错误，应为：时间窗口秒数,最大请求数，如：60,10")
		}

		var windowSeconds, maxRequests int
		_, err := fmt.Sscanf(parts[0], "%d", &windowSeconds)
		if err != nil || windowSeconds <= 0 {
			log.Fatal("时间窗口秒数必须是正整数")
		}

		_, err = fmt.Sscanf(parts[1], "%d", &maxRequests)
		if err != nil || maxRequests <= 0 {
			log.Fatal("最大请求数必须是正整数")
		}

		urlRateLimiter = NewURLRateLimiter(time.Duration(windowSeconds)*time.Second, maxRequests)
		log.Printf("URL限流配置：每个路径%d秒内最多%d次请求", windowSeconds, maxRequests)
	} else {
		// 默认限流配置：60秒内60次请求
		urlRateLimiter = NewURLRateLimiter(60*time.Second, 60)
		log.Printf("使用默认URL限流配置：每个路径60秒内最多60次请求")
	}

	proxy := &Proxy{
		requestBlock:   requestBlock,
		responseBlock:  responseBlock,
		requestKey:     []byte(requestKeyStr),
		responseKey:    []byte(responseKeyStr),
		address:        targetAddr,
		dataKey:        dataKey,
		client:         client,
		urlRateLimiter: urlRateLimiter,
	}

	// 打印配置信息
	log.Printf("监听端口: %s", port)
	log.Printf("请求解密密钥: %s", requestKeyStr)
	log.Printf("响应加密密钥: %s", responseKeyStr)
	log.Printf("目标地址: %s", u.String())
	log.Printf("数据键名: %s", dataKey)
	log.Printf("超时时间: %d秒", timeout)

	// 创建HTTP服务器
	server := &http.Server{
		Addr:         ":" + port,
		Handler:      http.HandlerFunc(proxy.handleRequest),
		ReadTimeout:  30 * time.Second,
		WriteTimeout: 30 * time.Second,
		IdleTimeout:  60 * time.Second,
	}

	log.Printf("代理服务器启动在端口 %s", port)
	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatal("启动服务失败:", err)
	}
}

// 标准化目标地址
func normalizeTargetAddress(targetAddr string) string {
	if !strings.Contains(targetAddr, ":") && len(targetAddr) < 5 {
		targetAddr = "127.0.0.1:" + targetAddr
	}
	if !strings.HasPrefix(targetAddr, "http") {
		targetAddr = "http://" + targetAddr
	}
	return targetAddr
}

func (p *Proxy) handleRequest(w http.ResponseWriter, r *http.Request) {
	// 添加请求ID用于日志追踪
	requestID := fmt.Sprintf("%d", time.Now().UnixNano())
	log.Printf("[%s] 收到请求: %s %s", requestID, r.Method, r.URL.Path)

	// 处理统计信息请求
	if r.URL.Path == "/_stats" {
		p.handleStats(w, r)
		return
	}

	// 限流检查（基于URL路径）
	if !p.urlRateLimiter.Allow(r.URL.Path) {
		current, max := p.urlRateLimiter.GetStats(r.URL.Path)
		log.Printf("[%s] 请求被限流，路径: %s, 当前窗口内请求数: %d/%d", requestID, r.URL.Path, current, max)
		w.Header().Set("X-RateLimit-Limit", fmt.Sprintf("%d", max))
		w.Header().Set("X-RateLimit-Remaining", "0")
		w.Header().Set("X-RateLimit-Reset", fmt.Sprintf("%d", time.Now().Add(p.urlRateLimiter.windowSize).Unix()))
		http.Error(w, "请求过于频繁，请稍后再试", http.StatusTooManyRequests)
		return
	}

	// 设置响应头
	w.Header().Set("X-Request-ID", requestID)
	w.Header().Set("X-Proxy-Server", "SM4-Proxy")

	// 设置限流响应头
	current, max := p.urlRateLimiter.GetStats(r.URL.Path)
	w.Header().Set("X-RateLimit-Limit", fmt.Sprintf("%d", max))
	w.Header().Set("X-RateLimit-Remaining", fmt.Sprintf("%d", max-current))

	// 处理请求
	var encryptedPayload []byte
	var err error

	switch r.Method {
	case http.MethodGet, http.MethodDelete:
		encryptedPayload, err = p.getQueryParamAsHex(r)
	case http.MethodPost, http.MethodPut, http.MethodPatch:
		encryptedPayload, err = p.readRequestBody(r)
	default:
		http.Error(w, "不支持的方法", http.StatusMethodNotAllowed)
		return
	}

	if err != nil {
		log.Printf("[%s] 解析请求失败: %v", requestID, err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// 解密请求数据
	plaintext, err := p.decrypt(encryptedPayload)
	if err != nil {
		log.Printf("[%s] 解密失败: %v", requestID, err)
		http.Error(w, "解密失败", http.StatusInternalServerError)
		return
	}

	// 转发请求
	resp, err := p.forwardRequest(r, plaintext, requestID)
	if err != nil {
		log.Printf("[%s] 转发请求失败: %v", requestID, err)
		http.Error(w, "转发请求失败", http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	// 读取响应
	rawResponseBody, err := io.ReadAll(resp.Body)
	if err != nil {
		log.Printf("[%s] 读取响应失败: %v", requestID, err)
		http.Error(w, "读取响应失败", http.StatusInternalServerError)
		return
	}

	// 加密响应
	encryptedResponseBody, err := p.encrypt(rawResponseBody)
	if err != nil {
		log.Printf("[%s] 加密响应失败: %v", requestID, err)
		http.Error(w, "加密响应失败", http.StatusInternalServerError)
		return
	}

	// 将加密后的响应包装在JSON对象中
	encryptedHex := hex.EncodeToString(encryptedResponseBody)
	responseData := map[string]string{
		p.dataKey: encryptedHex,
	}

	responseJSON, err := json.Marshal(responseData)
	if err != nil {
		log.Printf("[%s] JSON序列化失败: %v", requestID, err)
		http.Error(w, "JSON序列化失败", http.StatusInternalServerError)
		return
	}

	// 复制响应头
	for key, values := range resp.Header {
		for _, value := range values {
			w.Header().Add(key, value)
		}
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(resp.StatusCode)

	if _, err := w.Write(responseJSON); err != nil {
		log.Printf("[%s] 写入响应失败: %v", requestID, err)
	}

	log.Printf("[%s] 请求处理完成，状态码: %d", requestID, resp.StatusCode)
}

func (p *Proxy) getQueryParamAsHex(r *http.Request) ([]byte, error) {
	values := r.URL.Query()
	val, ok := values[p.dataKey]
	if !ok || len(val) == 0 {
		return nil, fmt.Errorf("GET请求中缺少'%s'参数", p.dataKey)
	}

	// 验证hex字符串格式
	if len(val[0])%2 != 0 {
		return nil, fmt.Errorf("无效的hex字符串长度")
	}

	return hex.DecodeString(val[0])
}

func (p *Proxy) readRequestBody(r *http.Request) ([]byte, error) {
	body, err := io.ReadAll(r.Body)
	if err != nil {
		return nil, fmt.Errorf("读取请求体失败: %v", err)
	}
	defer r.Body.Close()

	if len(body) == 0 {
		return nil, fmt.Errorf("POST请求中缺少请求体")
	}

	var jsonData map[string]interface{}
	if err := json.Unmarshal(body, &jsonData); err != nil {
		return nil, fmt.Errorf("解析JSON失败: %v", err)
	}

	dataValue, ok := jsonData[p.dataKey]
	if !ok {
		return nil, fmt.Errorf("JSON中缺少'%s'字段", p.dataKey)
	}

	dataStr, ok := dataValue.(string)
	if !ok {
		return nil, fmt.Errorf("'%s'字段必须是字符串类型", p.dataKey)
	}

	if dataStr == "" {
		return nil, fmt.Errorf("'%s'字段不能为空", p.dataKey)
	}

	// 验证hex字符串格式
	if len(dataStr)%2 != 0 {
		return nil, fmt.Errorf("无效的hex字符串长度")
	}

	return hex.DecodeString(dataStr)
}

func (p *Proxy) forwardRequest(r *http.Request, plaintext []byte, requestID string) (*http.Response, error) {
	// 创建新的请求
	var req *http.Request
	var err error

	if r.Method == http.MethodGet {
		// 智能处理解密后的数据
		queryParams, err := p.parseDecryptedData(plaintext)
		if err != nil {
			return nil, fmt.Errorf("解析解密数据失败: %v", err)
		}

		// 构建目标URL，但排除原始的dataKey参数
		parsedURL, err := url.Parse(p.address + r.URL.Path)
		if err != nil {
			return nil, fmt.Errorf("解析目标URL失败: %v", err)
		}

		// 获取原始查询参数，但排除dataKey
		originalQuery := r.URL.Query()
		originalQuery.Del(p.dataKey) // 删除原始的dataKey参数

		// 合并解密后的查询参数，也要排除dataKey
		for key, values := range queryParams {
			if key != p.dataKey { // 排除加密参数
				originalQuery[key] = append(originalQuery[key], values...)
			}
		}

		parsedURL.RawQuery = originalQuery.Encode()
		targetURL := parsedURL.String()

		req, err = http.NewRequestWithContext(r.Context(), r.Method, targetURL, nil)
	} else {
		// 对于其他方法，将解密后的数据作为请求体
		// 构建目标URL，但排除原始的dataKey参数
		parsedURL, err := url.Parse(p.address + r.URL.Path)
		if err != nil {
			return nil, fmt.Errorf("解析目标URL失败: %v", err)
		}

		// 获取原始查询参数，但排除dataKey
		originalQuery := r.URL.Query()
		originalQuery.Del(p.dataKey) // 删除原始的dataKey参数
		parsedURL.RawQuery = originalQuery.Encode()
		targetURL := parsedURL.String()

		req, err = http.NewRequestWithContext(r.Context(), r.Method, targetURL, bytes.NewReader(plaintext))
	}

	if err != nil {
		return nil, fmt.Errorf("创建请求失败: %v", err)
	}

	// 复制重要的请求头
	for key, values := range r.Header {
		if !strings.EqualFold(key, "Host") && !strings.EqualFold(key, "Content-Length") {
			for _, value := range values {
				req.Header.Add(key, value)
			}
		}
	}

	// 设置内容类型
	if r.Method != http.MethodGet {
		req.Header.Set("Content-Type", "application/json")
	}

	log.Printf("[%s] 转发请求到: %s", requestID, req.URL.String())

	// 发送请求
	return p.client.Do(req)
}

// parseDecryptedData 智能解析解密后的数据
// 支持两种格式：
// 1. 标准查询参数格式: key1=value1&key2=value2
// 2. JSON格式: {"key1":"value1","key2":"value2"}
func (p *Proxy) parseDecryptedData(plaintext []byte) (url.Values, error) {
	plaintextStr := strings.TrimSpace(string(plaintext))

	// 尝试解析为JSON
	if strings.HasPrefix(plaintextStr, "{") && strings.HasSuffix(plaintextStr, "}") {
		var jsonData map[string]interface{}
		if err := json.Unmarshal(plaintext, &jsonData); err == nil {
			// JSON解析成功，转换为查询参数
			queryParams := make(url.Values)
			for key, value := range jsonData {
				// 将各种类型的值转换为字符串
				switch v := value.(type) {
				case string:
					queryParams.Set(key, v)
				case float64:
					queryParams.Set(key, fmt.Sprintf("%.0f", v))
				case int:
					queryParams.Set(key, fmt.Sprintf("%d", v))
				case bool:
					queryParams.Set(key, fmt.Sprintf("%t", v))
				default:
					queryParams.Set(key, fmt.Sprintf("%v", v))
				}
			}
			log.Printf("JSON数据转换为查询参数: %v", queryParams)
			return queryParams, nil
		}
		// JSON解析失败，继续尝试标准查询参数格式
	}

	// 尝试解析为标准查询参数格式
	queryParams, err := url.ParseQuery(plaintextStr)
	if err != nil {
		return nil, fmt.Errorf("无法解析为查询参数或JSON格式: %v", err)
	}

	return queryParams, nil
}

func (p *Proxy) encrypt(plaintext []byte) ([]byte, error) {
	p.mu.RLock()
	defer p.mu.RUnlock()

	// 填充
	paddedData := pkcs5Pad(plaintext, p.responseBlock.BlockSize())

	ciphertext := make([]byte, len(paddedData))
	for i := 0; i < len(paddedData); i += p.responseBlock.BlockSize() {
		p.responseBlock.Encrypt(ciphertext[i:], paddedData[i:])
	}

	return ciphertext, nil
}

func (p *Proxy) decrypt(ciphertext []byte) ([]byte, error) {
	p.mu.RLock()
	defer p.mu.RUnlock()

	if len(ciphertext)%p.requestBlock.BlockSize() != 0 {
		return nil, errors.New("密文长度不是块大小的倍数")
	}

	plaintext := make([]byte, len(ciphertext))
	for i := 0; i < len(ciphertext); i += p.requestBlock.BlockSize() {
		p.requestBlock.Decrypt(plaintext[i:], ciphertext[i:])
	}

	// 去除填充
	unpaddedData, err := pkcs5Unpad(plaintext)
	if err != nil {
		return nil, fmt.Errorf("去除填充失败: %v", err)
	}

	return unpaddedData, nil
}

// PKCS#5填充
func pkcs5Pad(src []byte, blockSize int) []byte {
	padLen := blockSize - len(src)%blockSize
	padText := bytes.Repeat([]byte{byte(padLen)}, padLen)
	return append(src, padText...)
}

// PKCS#5去填充，增加错误检查
func pkcs5Unpad(src []byte) ([]byte, error) {
	if len(src) == 0 {
		return nil, errors.New("输入数据为空")
	}

	padLen := int(src[len(src)-1])
	if padLen > len(src) || padLen == 0 {
		return nil, errors.New("无效的填充长度")
	}

	// 验证填充
	for i := len(src) - padLen; i < len(src); i++ {
		if src[i] != byte(padLen) {
			return nil, errors.New("填充验证失败")
		}
	}

	return src[:len(src)-padLen], nil
}

// handleStats 处理统计信息请求
func (p *Proxy) handleStats(w http.ResponseWriter, r *http.Request) {
	// 获取所有URL路径的统计信息
	allStats := p.urlRateLimiter.GetAllStats()

	// 添加总体统计信息
	totalStats := map[string]interface{}{
		"url_stats":   allStats,
		"total_urls":  len(allStats),
		"window_size": p.urlRateLimiter.windowSize.Seconds(),
		"timestamp":   time.Now().Unix(),
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(totalStats)
}
